#include <cstdio>
#include <cstdlib>
//#include <tchar.h>

#include "opencv2/imgproc/imgproc.hpp"
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/video/video.hpp"
#include <opencv2/gpu/gpu.hpp>

#include "LcBasic.h"


#include "LcGraphCut.h"




class LcInterface
{
public:

	LcInterface();
	void input( int argc, char * const argv[]);

	//LcClassifier*  getClassifier(){ return my_classifier;}

	//LcVideoReadExt getTrainVideo();
	//LcVideoReadExt getTestVideo();


	bool IF_TRAIN;
	bool IF_TEST;

	int train_f_start;
	int train_f_end;
	int train_f_rate;

	int test_f_start;
	int test_f_end;
	int test_f_rate;

	string root;

	string train_video_name;
	string test_video_name;

	string feature_code_string;

	string classifier_name;

	void buildDatasetmap();

	map< string , Vec3i> data_map;

	bool setVideo( string name ,int & start, int & end , int & rate);
	void setClassifier();

	void init();
	
	int jac_down_rate;
	double learning_rate;
	double cooling_rate;
	int iter_times;

	string addition_string;
	//LcClassifier* my_classifier;
	


};

//LcVideoReadExt LcInterface::getTrainVideo()
//{
//	LcVideoReadExt ans;
//	ans.video_name = train_video_name;
//	ans.f_start = train_f_start;
//	ans.f_end = train_f_end;
//	ans.f_rate = train_f_rate;
//	ans.t_win = 0;
//	ans.root = root;
//	return ans;
//}
//
//LcVideoReadExt LcInterface::getTestVideo()
//{
//	LcVideoReadExt ans;
//	ans.video_name = test_video_name;
//	ans.f_start = test_f_start;
//	ans.f_end = test_f_end;
//	ans.f_rate = test_f_rate;
//	ans.t_win = 0;
//	ans.root = root;
//	return ans;
//}

void LcInterface::setClassifier()
{
	if(classifier_name == "rdt")
	{
		//my_classifier = new LcRandomTrees;
	}
}
void LcInterface::init()
{
	if(IF_TEST == false)
	{
		if(train_f_start<0) train_f_start = test_f_start;
		if(train_f_end<0) train_f_end = test_f_end;
		if(train_f_rate<0) train_f_rate = test_f_rate;
	}

	if(feature_code_string == "") feature_code_string = "l";

	if(classifier_name == "" ) classifier_name = "rdt";
	setClassifier();

	if(train_video_name == "" ){train_video_name = "GOPR2382_cab";}
	setVideo( train_video_name, train_f_start, train_f_end, train_f_rate);

	if(test_video_name == "" ){test_video_name = "GOPR2384_cab";}
	setVideo( test_video_name, test_f_start, test_f_end, test_f_rate);
}

LcInterface::LcInterface()
{
	train_f_start = -1;
	train_f_rate = -1;
	train_f_end = -1;

	test_f_start = -1;
	test_f_rate = -1;
	test_f_end = -1;

	root = "./../_FGS/";
	train_video_name = "";
	test_video_name = "";
	feature_code_string = "";
	classifier_name = "";

	IF_TEST = false;
	IF_TRAIN = false;

	buildDatasetmap();

	jac_down_rate = 10;
	learning_rate = 0.02;
	cooling_rate = 0;
	iter_times = 100;

	addition_string = "";
}

void LcInterface::buildDatasetmap()
{
	data_map.clear();
	data_map.insert( pair< string, Vec3i>( "GOPR2382_cab", Vec3i(150,5300,50)) );
	data_map.insert( pair< string, Vec3i>( "GOPR2384_cab", Vec3i(150,11200,25)) );
}

bool LcInterface::setVideo( string name ,int & start, int & end , int & rate)
{
	map<string , Vec3i>::iterator iter = data_map.find( name );
	if( iter != data_map.end())
	{
		Vec3i temp = iter->second;
		if(start<0)start = temp[0];
		if(end<0)end = temp[1];
		if(rate<0)rate = temp[2];
		return true;
	}
	else
	{
		cout << "One of Video is not in default" << endl;
		return false;
	}
}

void LcInterface::input( int argc, char * const argv[])
{
	for(int ind = 1;ind < argc;ind++)
	{
		/*** Parse switches ***/
		if (argv[ind][0] == '-') {
			switch (argv[ind][1]) {
				case 'f':
					root = string(argv[++ind]);
					break;
				case 't':
					train_video_name = string(argv[++ind]);
					break;
				case 'p':
					test_video_name = string(argv[++ind]);
					break;
				case 'v':
					feature_code_string =  string(argv[++ind]);
					break;
				case 'a':
					IF_TRAIN = true;
					break;
				case 'b':
					IF_TEST = true;
					break;
				case 's':
					//only for predict video setting , for trainning video it's always solid with frame
					// unless it's no testing model then the test frame setting will be use
					// as training parameter
					test_f_start =  atoi(argv[++ind]);
					break;
				case 'e':
					test_f_end = atoi(argv[++ind]);
					break;
				case 'r':
					test_f_rate = atoi(argv[++ind]);
					break;
				case 'c':
					classifier_name = string(argv[++ind]);
					break;
				case 'i':
					jac_down_rate = atoi(argv[++ind]);
					break;
				case 'j':
					learning_rate = atof(argv[++ind]);
					break;
				case 'k':
					iter_times = atoi(argv[++ind]);
					break;
				case 'l':
					cooling_rate = atof(argv[++ind]);
					break;
				case 'z':
					addition_string = string(argv[++ind]);
					break;


				default : printf("Unknown switch '%c'\n", argv[ind][1]);
				break;
			}
		}
	}

	init();
}


void getMainData(	LcInterface & my_setting, 
					vector<Mat> & node_descs,
					vector<Mat> & edge_descs,
					vector<Mat> & edge_maps,
					vector<Mat> & labels,
					vector<Mat> & label_maps,
					vector<Mat> & gts);


int main(int argc, char * const argv[])
{

	LcInterface my_setting;

	my_setting.input( argc ,argv);

	//char *const setting_string[] ={ "","-a","-t","GOPR2382_cab","-s", "550", "-e", "2000", "-r","50","-p","GOPR2382_cab","-f","./../../_FGS/","-i","1"};
	//char *const setting_string[] ={ "","-b","-t","GOPR2393_cab","-s", "200", "-e", "3000", "-r","400","-p","GOPR2393_cab","-f","./../../_FGS/","-i","1"};
	//my_setting.input( 15 + 1 ,setting_string);

	vector<Mat> node_descs;
	vector<Mat> edge_descs;
	vector<Mat> edge_maps;	
	vector<Mat> labels;
	vector<Mat> label_maps;
	vector<Mat> gts;

	getMainData( my_setting, node_descs, edge_descs , edge_maps, labels, label_maps, gts);

	LcMRF my_mrf;

	vector<int> dim;
	dim.push_back( node_descs[0].cols );
	dim.push_back( edge_descs[0].cols );

	my_mrf.set( dim );

	if(my_setting.IF_TRAIN)
	{
		Mat parameter;

		my_mrf.getParameter( parameter );		

		LcLMOpt my_Opt( node_descs,edge_descs, edge_maps, labels);

		my_Opt.setMRF( &my_mrf );

		my_Opt.iter_times = my_setting.iter_times;
		my_Opt.learning_rate = my_setting.learning_rate;
		my_Opt.cooling_rate = my_setting.cooling_rate;
		my_Opt.jac_down_rate = my_setting.jac_down_rate;

		my_Opt.work();

		parameter = my_Opt.my_parameter;

		{
			//save it
			string saving_name = my_setting.root + "MRF_output/para_"
				+ my_setting.train_video_name 
				+ my_setting.addition_string
				+ ".bin";

			lc::LcMat2Bin( saving_name.c_str() , parameter );
		}
	}
	else if(my_setting.IF_TEST)
	{
		Mat parameter;

		{
			//save it
			string saving_name = my_setting.root + "MRF_output/para_"
				+ my_setting.train_video_name 
				+ my_setting.addition_string
				+ ".bin";

			lc::LcBin2Mat( saving_name.c_str() , parameter );
		}

		my_mrf.setParameter( parameter );	

		{
			LcValidator _validator(0,0,0,0);

			LcValidator _pixel_validator(0,0,0,0);

			for(int i = 0 ; i < (int)labels.size();i++)
			{
				Mat res;

				_validator = _validator +
					my_mrf.predict( node_descs[i],
									edge_descs[i],
									edge_maps[i],
									res,
									labels[i]);

				Mat dst;
				superResReconstruct( res, label_maps[i],dst);

				Mat gt = gts[i];
				if(gt.type() == CV_8U);
				{
					gt.convertTo( gt, 5 ) ; gt = gt * (1.0/255.0);
				}

				_pixel_validator = _pixel_validator + LcValidator( dst, gt );

				
			}

			{
				string report_name = "./report/MRF_report_" + my_setting.train_video_name + "_on_"
					+ my_setting.test_video_name + my_setting.addition_string + ".txt";
				FILE * fout = fopen( report_name.c_str(),"w");				

				fprintf(fout,"%f\t%f\t%f\t%f\t%f\n", _validator.getF1(),
					_validator.tp,_validator.fp,_validator.fn,_validator.tn);

				fprintf(fout,"%f\t%f\t%f\t%f\t%f\n", _pixel_validator.getF1(),
					_pixel_validator.tp,_pixel_validator.fp,_pixel_validator.fn,_pixel_validator.tn);

				fclose(fout);
			}
		}
		
	}

	system("pause");

	return 0;
}




void getMainData(	LcInterface & my_setting, 
					vector<Mat> & node_descs,
					vector<Mat> & edge_descs,
					vector<Mat> & edge_maps,
					vector<Mat> & labels,
					vector<Mat> & label_maps,
					vector<Mat> & gts)
{
	bool IF_SAVE = false;		

	string root = my_setting.root;//"./../../_FGS/";
	string video_name = my_setting.train_video_name;//"GOPR2393_cab";
	

	int start_frame= my_setting.train_f_start;

	int end_frame = my_setting.train_f_end;

	int rate_frame = my_setting.train_f_rate;

	if(my_setting.IF_TEST)
	{
		video_name = my_setting.test_video_name;

		start_frame = my_setting.test_f_start;
		end_frame = my_setting.test_f_end;
		rate_frame = my_setting.test_f_rate;
	}

	string src_video = root + "vid/"+ video_name+".avi";

	VideoCapture src_cap( src_video );

	

	{
		int _n = (end_frame - start_frame)/rate_frame+5;

		node_descs.clear(); node_descs.reserve( _n );
		edge_descs.clear(); edge_descs.reserve( _n );
		 edge_maps.clear();  edge_maps.reserve( _n );
			labels.clear();     labels.reserve( _n );
		label_maps.clear();	label_maps.reserve( _n );
			   gts.clear();        gts.reserve( _n );
	}

	

	for(int i = 0 ; i< start_frame;i++){ Mat temp;src_cap >> temp;}

	for(int i = start_frame;i<=end_frame;i++)
	{
		Mat likelihood, src;

		src_cap>>src;

		if( (i-start_frame)%rate_frame!=0) continue;

		if(! src.data) break;		

		{
			char likeli_name[400];
			//sprintf( likeli_name , "./likelihood/frm_%s_%d.jpg", video_name.c_str(),i );
			sprintf( likeli_name , "../SparseTraining/likelihood/frm_%s_%d.jpg", video_name.c_str(),i );		
			likelihood = imread(likeli_name);
		}

		if(!likelihood.data) continue;

		cv::resize( src, src, likelihood.size() );

		if(likelihood.type() == CV_8UC3)
			cv::cvtColor( likelihood, likelihood, CV_RGB2GRAY);		

		Mat gt;

		{
			string msktyp = "msk_fg/";

			stringstream ss;
			ss << root + msktyp + video_name + "/" + video_name + "_maskgt_" << i <<".jpg";

			gt = imread(ss.str(),1);
		}

		if(!gt.data) continue;

		cv::resize( gt, gt, likelihood.size() );

		cout << "=============" << i << "=================" << endl;

		LcGraphCut my_graphcut;		

		my_graphcut.set( src, likelihood);

		

		my_graphcut.collectNodeStatistic( src, likelihood );

		{//collect Pb
			Mat Pb_result;

			char Pb_name[400];
			sprintf( Pb_name , "%s/Pb_output/frm_%s_%d.jpg", 
				my_setting.root.c_str(),
				video_name.c_str(),i );

			Pb_result = imread(Pb_name);
			
			

			if(Pb_result.data)
			{				
				if(Pb_result.type()==CV_8UC3)
					cv::cvtColor( Pb_result, Pb_result, CV_RGB2GRAY );

				cv::resize( Pb_result, Pb_result, likelihood.size() );

				my_graphcut.collectPb( Pb_result);
			}			
		}

		my_graphcut.collectHeatmap( likelihood );

		

		//my_graphcut.balanceBackGround( likelihood );

		Mat node_desc;

		{
			
			my_graphcut.computeAllNode( node_desc );						

			if(IF_SAVE)
			{
				char fname_node_desc[400];
				sprintf(fname_node_desc,"%s/superpixel_output/%s_node_desc_%d.bin"
					 ,root.c_str(), video_name.c_str(),i);
				lc::LcMat2Bin( fname_node_desc, node_desc);
			}

			if(0)
			{
				node_desc.convertTo(node_desc,5);
				cv::resize( node_desc, node_desc, Size(600,400));
				lc::colorshow("debug_node_desc",node_desc);
				cv::waitKey(1);
			}

		}

		Mat edge_desc;
		Mat edge_map;

		{
			
			
			my_graphcut.computeAllEdge( edge_map, edge_desc );

			if(IF_SAVE)
			{
				char fname_edge_desc[400];
				sprintf(fname_edge_desc,"%s/superpixel_output/%s_edge_desc_%d.bin"
					,root.c_str(), video_name.c_str(),i);
				lc::LcMat2Bin( fname_edge_desc, edge_desc);

				Mat _edge_map;
				edge_map.convertTo(_edge_map,5);
				char fname_edge_map[400];
				sprintf(fname_edge_map,"%s/superpixel_output/%s_edge_map_%d.bin"
					,root.c_str(), video_name.c_str(),i);
				lc::LcMat2Bin( fname_edge_map, _edge_map);
			}


			if(0)
			{
				edge_desc.convertTo(edge_desc,5);
				cv::resize( edge_desc, edge_desc, Size(600,400));
				lc::colorshow("debug_edge_desc",edge_desc);
				cv::waitKey(1);
			}
		}

		Mat label;

		{			
			my_graphcut.getLabel( gt, label);

			char fname_label[400];
			sprintf(fname_label,"%s/superpixel_output/%s_label_%d.bin"
				,root.c_str(), video_name.c_str(),i);
			lc::LcMat2Bin( fname_label, label);
		}

		if( label.rows > 0 && node_desc.rows > 0)
		{
			node_descs.push_back( node_desc );
			edge_descs.push_back( edge_desc );
			edge_maps.push_back(  edge_map  );
			labels.push_back(     label     );
			label_maps.push_back( my_graphcut.label_map );
			gts.push_back( gt );

		}
	}
}